{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "836ed7db",
   "metadata": {},
   "source": [
    "# Consuming Multiple Return Values from an Electron\n",
    "\n",
    "In Python, it is possible to return multiple values from a function. The technique explained here prevents unnecessary usage of compute resources when you consume the results of such a function in a Covalent workflow.\n",
    "\n",
    "## Context\n",
    "\n",
    "When a function consumes one of multiple output values of another function – that is, uses the value as an input — Python creates a new instance of the upstream function to feed the value to the downstream function. This consumes additional memory and CPU overhead. In Covalent the copied functions inherit their dependencies, including the executor, from the upstream electron. As a result, the duplicated electron can represent a hefty chunk of memory and other resources just to retrieve the indexed result value. This extraneous activity on a potentially high-value resource is completely avoidable. \n",
    "\n",
    "\n",
    "## Best Practice\n",
    "\n",
    "Instead of taking the upstream values one at a time, input them as a tuple to the downstream function. The function reads the tuple without the creation of new upstream electrons.\n",
    "\n",
    "## Example\n",
    "\n",
    "Contrast the two examples below. \n",
    "\n",
    "### Example 1: Not Recommended\n",
    "\n",
    "This example demonstrates a questionable approach: consuming multiple return values separately.\n",
    "\n",
    "The following simple workflow shows the creation of the extraneous electrons (see the *Node Outputs* in the Lattice Result). Nodes 3 and 5 are created only to supply the output variables from the computation in `task_1` (node 0). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "1b8fc55a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Lattice Result\n",
      "==============\n",
      "status: COMPLETED\n",
      "result: (16, 36)\n",
      "input args: ['2', '3']\n",
      "input kwargs: {}\n",
      "error: None\n",
      "\n",
      "start_time: 2023-03-13 19:46:14.283738\n",
      "end_time: 2023-03-13 19:46:14.801465\n",
      "\n",
      "results_dir: /Users/dave/.local/share/covalent/data\n",
      "dispatch_id: 8854822f-c28a-46f6-8df7-254b2ee31239\n",
      "\n",
      "Node Outputs\n",
      "------------\n",
      "task_1(0): (4, 6)\n",
      ":parameter:2(1): 2\n",
      ":parameter:3(2): 3\n",
      ":task_1()[0](3): 4\n",
      ":parameter:0(4): 0\n",
      ":task_1()[1](5): 6\n",
      ":parameter:1(6): 1\n",
      "task_2(7): (16, 36)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import covalent as ct\n",
    "\n",
    "# Technique 1: Questionable\n",
    "\n",
    "@ct.electron\n",
    "def task_1(x, y):\n",
    "    return x * 2, y * 2\n",
    "\n",
    "@ct.electron\n",
    "def task_2(x, y):\n",
    "    return x ** 2, y ** 2\n",
    "\n",
    "@ct.lattice\n",
    "def workflow_1(a, b):\n",
    "    res_1, res_2 = task_1(a, b) # Multiple outputs\n",
    "    return task_2(res_1, res_2) # Not optimal: Consume the outputs individually\n",
    "\n",
    "dispatch_id_1 = ct.dispatch(workflow_1)(2,3)\n",
    "result_1 = ct.get_result(dispatch_id_1, wait=True)\n",
    "\n",
    "print(result_1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51550e5b",
   "metadata": {},
   "source": [
    "### Example 2: Improved\n",
    "\n",
    "In contrast, the following workflow demonstrates that, by consuming the outputs as a list, the workflow avoids creating unnecessary nodes. Notice that only one instance of `task_1` is created."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e5dc7373",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Lattice Result\n",
      "==============\n",
      "status: COMPLETED\n",
      "result: (16, 36)\n",
      "input args: ['2', '3']\n",
      "input kwargs: {}\n",
      "error: None\n",
      "\n",
      "start_time: 2023-03-13 19:46:16.505981\n",
      "end_time: 2023-03-13 19:46:16.762871\n",
      "\n",
      "results_dir: /Users/dave/.local/share/covalent/data\n",
      "dispatch_id: 63710c74-9313-4bda-ad47-bd88fc8591ac\n",
      "\n",
      "Node Outputs\n",
      "------------\n",
      "task_1(0): (4, 6)\n",
      ":parameter:2(1): 2\n",
      ":parameter:3(2): 3\n",
      "task_2_new(3): (16, 36)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import covalent as ct\n",
    "\n",
    "# Technique 2: Better\n",
    "\n",
    "# The same task_1 as in the previous example\n",
    "@ct.electron\n",
    "def task_1(x, y):\n",
    "    return x * 2, y * 2\n",
    "\n",
    "# New task_2 takes an array argument\n",
    "@ct.electron\n",
    "def task_2_new(arr):\n",
    "    return arr[0] ** 2, arr[1] ** 2\n",
    "\n",
    "\n",
    "@ct.lattice\n",
    "def workflow_2(a, b): \n",
    "    res = task_1(a, b)     # Capture multiple outputs in an array\n",
    "    return task_2_new(res) # Optimal: consume the array\n",
    "\n",
    "\n",
    "dispatch_id_2 = ct.dispatch(workflow_2)(2,3)\n",
    "result_2 = ct.get_result(dispatch_id_2, wait=True)\n",
    "\n",
    "print(result_2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00527e50",
   "metadata": {},
   "source": [
    "## See Also\n",
    "\n",
    "[Creating and Assigning an Executor](./executor_assignment.ipynb)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
